""" ocr web 服务 """
import json
import os
import socket
import tempfile
import threading
from pathlib import Path
from time import time
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from typing import Literal

from bottle import route, run, post, request, BaseRequest, response

from PPOCR_api import GetOcrApi
from RapidOCR_api import OcrAPI

ALLOWED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.pdf'}
IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg'}

# OCR配置

SERVER_DIR = Path(__file__).resolve().parent
Rapid_OCR_PATH = str(SERVER_DIR.joinpath(Path("OCRCore/RapidOCR-json_v0.2.0/RapidOCR-json.exe")))
Paddle_OCR_PATH = str(SERVER_DIR.joinpath(Path("OCRCore/PaddleOCR-json_v1.4.1_windows_x64/PaddleOCR-json_v1.4.1/PaddleOCR-json.exe")))
TEST_OCR_FILE_PATH = str(SERVER_DIR.joinpath(Path("images/test_ocr.png")))
max_workers = 24

# def ocr_core_factory(ocr_core_type: Literal["RapidOCR", "PaddleOCR"] = "RapidOCR"):
#     """
#     OCR核心工厂
#     """
#     if ocr_core_type == "PaddleOCR":
#         return PaddleOCR
#     elif ocr_core_type == "RapidOCR":
#         return RapidOCR
#     return None


def start_bottle(port: int = 9911, memory_size=300 * 1024 * 1024):
    """
    bottle 启动函数
    :param port:
    :param memory_size: 300MB
    :return:
    """
    # 设置文件大小限制
    BaseRequest.MEMFILE_MAX = memory_size  # 300MB
    run(port=port)


def is_port_in_use(port):
    """
    检查端口是否被占用
    :param port:
    :return:
    """
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) == 0


def start_bottle_thread(port: int = 9911):
    """
    启动bottle的线程
    :param port:
    :return:
    """
    threading.Thread(target=start_bottle, args=(port,), daemon=True).start()


def process_image_memory_task(file_bytes,
                              ocr_core_type: Literal["RapidOCR", "PaddleOCR"] = "RapidOCR"):
    """
    处理内存中的图片任务
    """
    from PPOCR_api import GetOcrApi
    from RapidOCR_api import OcrAPI
    if ocr_core_type == "PaddleOCR":
        ocr_core = GetOcrApi(Paddle_OCR_PATH)

        try:
            res = ocr_core.runBytes(file_bytes)
            return res
        finally:
            ocr_core.exit()
    elif ocr_core_type == "RapidOCR":
        ocr_core =  OcrAPI(Rapid_OCR_PATH)
        try:
            res = ocr_core.runBytes(file_bytes)
            return res
        finally:
            ocr_core.stop()
    return None


def process_image_file_task(image_path,
                            ocr_core_type: Literal["RapidOCR", "PaddleOCR"] = "RapidOCR"):
    """
    处理文件路径的图片任务
    """
    from PPOCR_api import GetOcrApi
    from RapidOCR_api import OcrAPI
    if ocr_core_type == "PaddleOCR":
        ocr_core = GetOcrApi(Paddle_OCR_PATH)
        try:
            res = ocr_core.run(image_path)
            return res
        finally:
            ocr_core.exit()
    elif ocr_core_type == "RapidOCR":
        ocr_core = OcrAPI(Rapid_OCR_PATH)
        try:
            res = ocr_core.run(image_path)
            return res
        finally:
            ocr_core.stop()
    return None

def pdf_to_images(pdf_path, output_dir, file_name):
    """
    pdf转图片
    :param pdf_path: PDF文件路径
    :param output_dir: 输出目录
    :param file_name: 文件名前缀
    :return: 图片文件路径列表
    """
    import fitz
    pdf_doc = fitz.open(pdf_path)
    file_names = []
    for page_index in range(pdf_doc.page_count):
        page = pdf_doc[page_index]
        zoom = (1.3, 1.3)
        mat = fitz.Matrix(zoom[0], zoom[1])
        pix = page.get_pixmap(matrix=mat, alpha=False)
        temp_file_name = file_name + str(page_index) + '.png'
        image_abs_path = os.path.join(output_dir, temp_file_name)
        pix.save(image_abs_path)
        file_names.append(image_abs_path)
    return file_names


def get_file_extension(filename):
    """
    获取文件扩展名
    :param filename:
    :return:
    """
    return os.path.splitext(filename)[1].lower()


def allowed_file(filename):
    """
    判断文件是否为允许的类型
    :param filename:
    :return:
    """
    return get_file_extension(filename) in ALLOWED_EXTENSIONS


@route('/ping')
def ping():
    """
    连通性测试
    :return:
    """
    return json.dumps({"code": 200, "msg": "pong"})


@route('/ocr_test')
def ocr_test():
    """
    测试ocr 是否可用
    :return:
    """
    try:
        # 支持的ocr
        supported_ocr = []
        for ocr_core_type in ["PaddleOCR", "RapidOCR"]:
            result = process_image_file_task(TEST_OCR_FILE_PATH,ocr_core_type=ocr_core_type)
            print(result)
            if result:
                supported_ocr.append(ocr_core_type)
        return json.dumps({"code": 200, "msg": "ocr 可用", "data": supported_ocr})
    except Exception as e:
        return json.dumps({"code": 500, "msg": f"ocr 不可用: {str(e)}"})


@post("/ocr_sse")
def ocr_api_sse():
    """
    SSE接口，用于流式返回OCR结果
    :return:
    """
    response.content_type = 'text/event-stream'
    response.cache_control = 'no-cache'

    def event_stream():
        """
        流式返回数据
        :return:
        """
        upload_files = request.files.getall('file')
        if not upload_files:
            yield f"data: {json.dumps({'code': 400, 'error': '请上传文件'})}\n\n"
            yield "event: end\ndata: {}\n\n"
            return

        start_time = time()
        tasks = []

        # 准备所有任务
        with tempfile.TemporaryDirectory() as tmp_dir:
            for upload_file in upload_files:
                if not allowed_file(upload_file.raw_filename):
                    yield f"data: {json.dumps({'filename': upload_file.raw_filename, 'error': f'不支持的文件类型: {get_file_extension(upload_file.raw_filename)}', 'status': 'error'})}\n\n"
                    continue

                file_extension = get_file_extension(upload_file.raw_filename)

                if file_extension in {'.png', '.jpg', '.jpeg'}:
                    file_bytes = upload_file.file.read()
                    tasks.append({
                        'type': 'image_memory',
                        'filename': upload_file.raw_filename,
                        'data': file_bytes
                    })
                elif file_extension == '.pdf':
                    yield f"data: {json.dumps({'filename': upload_file.raw_filename, 'status': 'processing_pdf', 'message': '正在转换PDF为图片'})}\n\n"
                    file_path = os.path.join(tmp_dir, upload_file.raw_filename)
                    upload_file.save(file_path)
                    image_files = pdf_to_images(file_path, tmp_dir, upload_file.raw_filename[:-4])
                    for i, image_path in enumerate(image_files):
                        tasks.append({
                            'type': 'image_file',
                            'filename': f"{upload_file.raw_filename}_page_{i + 1}",
                            'data': image_path,
                            'original_filename': upload_file.raw_filename,
                            'page_num': i + 1
                        })

            # 多线程处理任务并流式返回结果

            yield f"data: {json.dumps({'status': 'processing', 'total_tasks': len(tasks), 'workers': max_workers})}\n\n"

            with ProcessPoolExecutor(max_workers=max_workers) as executor:
                future_to_task = {}
                for task in tasks:
                    if task['type'] == 'image_memory':
                        future = executor.submit(process_image_memory_task, task['data'])
                    else:  # image_file
                        future = executor.submit(process_image_file_task, task['data'])
                    future_to_task[future] = task

                completed = 0
                for future in future_to_task:
                    task = future_to_task[future]
                    try:
                        result_data = future.result()
                        completed += 1
                        if task['type'] == 'image_memory':
                            yield f"data: {json.dumps({'filename': task['filename'], 'data': result_data, 'completed': completed, 'total': len(tasks)})}\n\n"
                        else:  # image_file (PDF页面)
                            yield f"data: {json.dumps({'filename': task['filename'], 'original_filename': task['original_filename'], 'page_num': task['page_num'], 'data': result_data, 'completed': completed, 'total': len(tasks)})}\n\n"
                    except Exception as e:
                        completed += 1
                        if task['type'] == 'image_memory':
                            yield f"data: {json.dumps({'filename': task['filename'], 'error': str(e), 'completed': completed, 'total': len(tasks)})}\n\n"
                        else:  # image_file
                            yield f"data: {json.dumps({'filename': task['filename'], 'original_filename': task['original_filename'], 'page_num': task['page_num'], 'error': str(e), 'completed': completed, 'total': len(tasks)})}\n\n"

        yield f"data: {json.dumps({'status': 'completed', 'total_time': time() - start_time})}\n\n"
        yield "event: end\ndata: {}\n\n"

    return event_stream()


@post("/ocr")
def ocr():
    """
    ocr 接口
    :return:
    """
    upload_files = request.files.getall('file')
    if not upload_files:
        return json.dumps({"code": 400, "msg": "请上传文件"})

    start_time = time()

    # 收集所有需要处理的任务
    tasks = []
    with tempfile.TemporaryDirectory() as tmp_dir:
        for upload_file in upload_files:
            if not allowed_file(upload_file.raw_filename):
                tasks.append({
                    'type': 'error',
                    'filename': upload_file.raw_filename,
                    'error': f"不支持的文件类型: {get_file_extension(upload_file.raw_filename)}"
                })
                continue

            file_extension = get_file_extension(upload_file.raw_filename)

            if file_extension in IMAGE_EXTENSIONS:
                # 图片文件，保存字节数据用于内存处理
                file_bytes = upload_file.file.read()
                tasks.append({
                    'type': 'image_memory',
                    'filename': upload_file.raw_filename,
                    'data': file_bytes
                })
            elif file_extension == '.pdf':
                # PDF 文件，需要转换为图片
                file_path = os.path.join(tmp_dir, upload_file.raw_filename)
                upload_file.save(file_path)
                image_files = pdf_to_images(file_path, tmp_dir, upload_file.raw_filename[:-4])
                for i, image_path in enumerate(image_files):
                    tasks.append({
                        'type': 'image_file',
                        'filename': f"{upload_file.raw_filename}_page_{i + 1}",
                        'data': image_path,
                        'original_filename': upload_file.raw_filename,
                        'page_num': i + 1
                    })

        # 使用多线程处理所有任务
        results = []
        pdf_grouped_results = {}

        print(f"开始处理 {len(tasks)} 个任务，使用 {max_workers} 个线程")

        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有任务
            future_to_task = {}
            for task in tasks:
                if task['type'] == 'error':
                    # 错误任务直接添加到结果中
                    results.append({
                        "filename": task['filename'],
                        "error": task['error']
                    })
                elif task['type'] == 'image_memory':
                    future = executor.submit(process_image_memory_task, task['data'])
                    future_to_task[future] = task
                elif task['type'] == 'image_file':
                    future = executor.submit(process_image_file_task, task['data'])
                    future_to_task[future] = task

            # 收集所有结果
            for future in future_to_task:
                task = future_to_task[future]
                try:
                    result_data = future.result()
                    if task['type'] == 'image_memory':
                        results.append({
                            "filename": task['filename'],
                            "data": result_data
                        })
                    elif task['type'] == 'image_file':
                        # 处理PDF页面结果
                        original_filename = task['original_filename']
                        if original_filename not in pdf_grouped_results:
                            pdf_grouped_results[original_filename] = []
                        pdf_grouped_results[original_filename].append({
                            "page_num": task['page_num'],
                            "data": result_data
                        })
                except Exception as e:
                    if task['type'] == 'image_memory':
                        results.append({
                            "filename": task['filename'],
                            "error": str(e)
                        })
                    elif task['type'] == 'image_file':
                        original_filename = task['original_filename']
                        if original_filename not in pdf_grouped_results:
                            pdf_grouped_results[original_filename] = []
                        pdf_grouped_results[original_filename].append({
                            "page_num": task['page_num'],
                            "error": str(e)
                        })

        # 整理PDF结果
        for original_filename, page_results in pdf_grouped_results.items():
            # 按页码排序
            page_results.sort(key=lambda x: x.get('page_num', 0))
            results.append({
                "filename": original_filename,
                "data": page_results
            })

    print("多线程处理总耗时：", (time() - start_time))
    return json.dumps({"code": 200, "msg": "success", "data": results})


if __name__ == '__main__':
    start_bottle()
